#include "SubProblem.hpp"
#include <cstring>
#include <iostream>

using namespace std;
using namespace Ipopt;

#define COMPUTE_NEW_ARG \
	if (new_z) { \
		computeNewArg(z); \
	}

SubProblem::SubProblem(AlgorithmRef algEngine, Index procIndex,
		Index startIndex, Index endIndex) :
	algEngine(algEngine), startIndex(startIndex), endIndex(endIndex),
			procIndex(procIndex) {
	originalDim = algEngine.dim;
	subspaceDim = endIndex - startIndex + algEngine.subProblemsCount - 1;
	procCount = algEngine.subProblemsCount;

	if (!allocStructures()) {
		//TODO throw
		assert(false);
	}

	//dane jakobianu
	subOptNnzJac = algEngine.procToJacElemNum[procIndex];
	subNnzJac = subOptNnzJac + algEngine.procToJacBlockNumDist[procCount]
			- algEngine.procToJacBlockNum[procIndex];
}

/** default destructor */
SubProblem::~SubProblem() {
	deallocStructures();
}

/** Method to return some info about the nlp
 * n - rozmiar problemu
 * m - ilosc ograniczen
 * nnz_jac_g - number of non-zero entries in jacobian
 * nnz_h_lag - number of non-zero entries in hessian
 * index_style - C_STYLE, numbering used for adressing rows/cols in sparse matrix
 *
 * */
bool SubProblem::get_nlp_info(Index& n, Index& m, Index& nnz_jac_g,
		Index& nnz_h_lag, IndexStyleEnum& index_style) {

	n = subspaceDim;
	m = algEngine.restrictions;
	nnz_jac_g = subNnzJac;
	nnz_h_lag = 0;
	index_style = C_STYLE;

	return true;
}

/** Method to return the bounds for my problem
 * n  -rozmiar problemu
 * x_l, x_u - ograniczenia na x
 * m - ilosc ograniczen
 * g_l,g_u, ograniczenia na g(x); nlp_lower_bound_inf
 * */
bool SubProblem::get_bounds_info(Index n, Number* x_l, Number* x_u, Index m,
		Number* g_l, Number* g_u) {

	int ssIndex = 0;
	//subspace which is being optimized
	for (int i = startIndex; i < endIndex; ++i, ssIndex++) {
		x_l[ssIndex] = algEngine.x_l[i];
		x_u[ssIndex] = algEngine.x_u[i];
	}
	//can't constrain lambdas
	while (ssIndex < n) {
		x_l[ssIndex] = -1e19;
		x_u[ssIndex] = 1e19;
		++ssIndex;
	}

	if (m > 0) {
		//constraints
		for (int i = 0; i < m; ++i) {
			g_l[i] = algEngine.g_l[i];
			g_u[i] = algEngine.g_u[i];
		}
	}

	return true;
}

/** Method to return the starting point for the algorithm
 * init_x - if true, starting point must be provided
 * x - starting point
 * init_z - if true one must provide initial value for bound multipliers z_L, z_U
 * m - number of constraints
 * init_lambda - if true, one must provide initial values for constraint multipliers
 * lamda - constraint multipliers values
 * dla domyslnych opcji, init_z, init_lamda sa false (problem dualny)
 * */
bool SubProblem::get_starting_point(Index n, bool init_x, Number* x,
		bool init_z, Number* z_L, Number* z_U, Index m, bool init_lambda,
		Number* lambda) {
	assert(init_z == false);
	assert(init_lambda == false);

	if (init_x) {
		Index ssIndex = 0;

		//subspace which is being optimized
		for (int i = startIndex; i < endIndex; ++i, ssIndex++) {
			x[ssIndex] = algEngine.lastOptX[i];
		}

		//complement subspace
		while (ssIndex < subspaceDim) {
			x[ssIndex++] = 1;
		}
	}

	return true;
}

/** Method to return the objective value */
bool SubProblem::eval_f(Index n, const Number* z, bool new_z, Number& outValue) {
	assert(n == subspaceDim);

	COMPUTE_NEW_ARG;

	outValue = funValue;

	return true;
}

/** Method to return the gradient of the objective
 * n - number of variables
 * x - value at which gradient of f is to be evaluated
 * new_x - false if any eval method was previously called with the same values
 * grad_f - output gradien
 */
bool SubProblem::eval_grad_f(Index n, const Number* z, bool new_z,
		Number* outGradient) {
	assert(n == subspaceDim);

	COMPUTE_NEW_ARG;

	int ssIndex = 0;
	//subspace which is being optimized
	for (int i = startIndex; i < endIndex; ++i, ssIndex++) {
		outGradient[ssIndex] = funGradient[i];
	}

	//complement subspace
	for (int p = 0; p < procCount; ++p) {
		if (p != procIndex) {
			Number grad_p = 0;
			int endIndex = algEngine.procStartIndex[p + 1];
			for (int i = algEngine.procStartIndex[p]; i < endIndex; ++i) {
				grad_p += algEngine.optFunDiagMatrix[i] * funGradient[i];
			}
			outGradient[ssIndex++] = grad_p;
		}
	}

	return true;
}

/** Method to return the constraint residuals */
bool SubProblem::eval_g(Index n, const Number* z, bool new_z, Index m,
		Number* g) {
	COMPUTE_NEW_ARG;

	if (m > 0) {
		return algEngine.problemDefinition.eval_g(originalDim, xOrgSpace, true,
				m, g);
	}
	return false;
}

bool SubProblem::eval_jac_g(Index n, const Number* z, bool new_z, Index m,
		Index nele_jac, Index* iRow, Index *jCol, Number* values) {
	COMPUTE_NEW_ARG;

	if (m > 0) {
		if (values == NULL) {
			assert(subNnzJac == nele_jac);
			Index jacOptInd = 0;
			for (Index orgNzi = 0; orgNzi < algEngine.orgNnzJac; orgNzi++) {
				if (algEngine.jacJCol[orgNzi] >= startIndex
						&& algEngine.jacJCol[orgNzi] < endIndex) {
					//optimized subspace
					iRow[jacOptInd] = algEngine.jacIRow[orgNzi];
					jCol[jacOptInd] = algEngine.jacJCol[orgNzi] - startIndex;
					assert(jCol[jacOptInd] >= 0 && jCol[jacOptInd] < subspaceDim);
					assert(iRow[jacOptInd] >= 0 && iRow[jacOptInd] < m);
					jacOptInd++;
				} else {
					//complement subspace
					Index proc = algEngine.nnzJacIndToProc(orgNzi);
					Index
							jacSubInd =
									subOptNnzJac
											+ algEngine.procToJacBlockNumDist[proc]
											+ algEngine.restrToProcJacBlockInd[algEngine.jacIRow[orgNzi]];
					if (proc > procIndex) {
						jacSubInd -= algEngine.procToJacBlockNum[procIndex];
					}

					assert(jacSubInd < subNnzJac && jacSubInd >= endIndex - startIndex);
					iRow[jacSubInd] = algEngine.jacIRow[orgNzi];
					jCol[jacSubInd] = endIndex - startIndex + proc;
					if (proc > procIndex) {
						jCol[jacSubInd]--;
					}
					assert(jCol[jacSubInd] >= 0 && jCol[jacSubInd] < subspaceDim);
					assert(iRow[jacSubInd] >= 0 && iRow[jacSubInd] < m);
				}
			}
			assert(jacOptInd == subOptNnzJac);
		} else {
			//wartosci
			algEngine.problemDefinition.eval_jac_g(originalDim, xOrgSpace,
					true, m, algEngine.orgNnzJac, algEngine.jacIRow,
					algEngine.jacJCol, orgJacob);

			memset(values, 0, nele_jac * sizeof(Number));
			Index jacOptInd = 0;

			//iterujemy po elementach oryginalnego jakobianu, kazdy jakos wplywa na te w podproblemie
			for (Index orgNzi = 0; orgNzi < algEngine.orgNnzJac; orgNzi++) {
				if (algEngine.jacJCol[orgNzi] >= startIndex
						&& algEngine.jacJCol[orgNzi] < endIndex) {
					//optimized subspace
					values[jacOptInd++] = orgJacob[orgNzi];
				} else {
					//complement subspace
					Index proc = algEngine.nnzJacIndToProc(orgNzi);
					Index
							jacSubInd =
									subOptNnzJac
											+ algEngine.procToJacBlockNumDist[proc]
											+ algEngine.restrToProcJacBlockInd[algEngine.jacIRow[orgNzi]];
					if (proc > procIndex) {
						jacSubInd -= algEngine.procToJacBlockNum[procIndex];
					}

					assert(jacSubInd < subNnzJac && jacSubInd >= endIndex - startIndex);
					values[jacSubInd]
							+= orgJacob[orgNzi]
									* algEngine.optFunDiagMatrix[algEngine.jacJCol[orgNzi]];
				}
			}
			assert(jacOptInd == subOptNnzJac);
		}
	}
	return true;
}

/** Method to return:
 *   1) The structure of the hessian of the lagrangian (if "values" is NULL)
 *   2) The values of the hessian of the lagrangian (if "values" is not NULL)
 */
bool SubProblem::eval_h(Index n, const Number* z, bool new_z,
		Number obj_factor, Index m, const Number* lambda, bool new_lambda,
		Index nele_hess, Index* iRow, Index* jCol, Number* values) {
	COMPUTE_NEW_ARG;

	return false;
}

/** This method is called when the algorithm is complete so the TNLP can store/write the solution */
void SubProblem::finalize_solution(SolverReturn status, Index n,
		const Number* x, const Number* z_L, const Number* z_U, Index m,
		const Number* g, const Number* lambda, Number obj_value IPOPT8_FINARGS) {
	computeNewArg(x);
}

void SubProblem::computeNewArg(const Number *z) {
	//z is point in subspace

	int ssIndex = 0;

	//subspace which is being optimized
	for (int i = startIndex; i < endIndex; ++i, ++ssIndex) {
		xOrgSpace[i] = z[ssIndex];
	}

	//complement subspace
	for (int p = 0; p < procCount; ++p) {
		if (p != procIndex) {
			int endIndex = algEngine.procStartIndex[p + 1];
			for (int i = algEngine.procStartIndex[p]; i < endIndex; ++i) {
				xOrgSpace[i] = algEngine.lastOptX[i] + z[ssIndex]
						* algEngine.optFunDiagMatrix[i];
			}
			++ssIndex;
		}
	}

	//compute function value
	algEngine.problemDefinition.eval_f(originalDim, xOrgSpace, true, funValue);

	//compute function gradient
	algEngine.problemDefinition.eval_grad_f(originalDim, xOrgSpace, true,
			funGradient);
}

bool SubProblem::allocStructures() {
	//original function cache
	assert(originalDim > 0);
	Number* pData = new Number[2 * originalDim];
	orgJacob = new Number[algEngine.orgNnzJac];

	if (pData == NULL || orgJacob == NULL) {
		deallocStructures();
		return false;
	}

	xOrgSpace = pData;
	funGradient = xOrgSpace + originalDim;

	return true;
}

void SubProblem::deallocStructures() {
	if (xOrgSpace != NULL) {
		delete[] xOrgSpace;
		xOrgSpace = NULL;
		funGradient = NULL;
	}

	if (orgJacob != NULL) {
		delete[] orgJacob;
		orgJacob = NULL;
	}
}
